////////////////////////////////////////////////////////////////////////
//
// Assignment 4 - Car tracking
//
// The fourth assignment for Advanced Multimedia Computing
//
// T. Kostelijk & F. Huizinga
//     0418889      0418862
//
////////////////////////////////////////////////////////////////////////
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <cv.h>
#include <highgui.h>
#include <assert.h>

const char  *WINDOW        = "main";  /* Main window */
const char  *BGEST         = "bgest"; /* Window after background estimation */
const char  *MCONV         = "mconv"; /* Window after convolution */
const int    KERNEL_SIZE   = 5;       /* Size of the gaussian kernel */
const int    HISTORY       = 10;      /* Number of history frame (context) */
const int    SUPPORT       = 5;       /* Size of structuring element */
const int    THRESHOLD     = 180;     /* Threshold on normal image */
int          counter       = 0;

/* Save image to harddisk */
inline void save(char *name, IplImage *img) {
	char buf[256];
	sprintf(buf, "%8.8d-%s.png", counter, name);
	if (!cvSaveImage(buf, img))
		printf("Could not save %s\n", name);
	else
		counter++;
}

/* Load image from harddisk */
IplImage* loadImage(int index, int type, char *dir) {
	char filename[256];
	sprintf(filename, "%s%8.8d.png", dir, index);
	IplImage *img = cvLoadImage(filename, type);
	return img;
}

/* Perform bg estimation on HISTORY amount of images */
IplImage* bgEstimate(int index, CvSize &size, char *dir) {
	assert(index >= HISTORY);

	IplImage *img, *imgSmooth, *diff, *avg;
	imgSmooth = cvCreateImage(size, IPL_DEPTH_8U, 1);
	diff      = cvCreateImage(size, IPL_DEPTH_8U, 1);
	avg       = cvCreateImage(size, IPL_DEPTH_8U, 1);

	for (int i = index-HISTORY; i < index; i++) {
		img = loadImage(i, CV_LOAD_IMAGE_GRAYSCALE, dir);
		cvSmooth(img, imgSmooth, CV_GAUSSIAN, KERNEL_SIZE, KERNEL_SIZE);
		cvConvertScale(imgSmooth, imgSmooth, 1.0/HISTORY);
		cvAdd(avg, imgSmooth, avg);
	}
	img = loadImage(index, CV_LOAD_IMAGE_GRAYSCALE, dir);
	cvSmooth(img, imgSmooth, CV_GAUSSIAN, KERNEL_SIZE, KERNEL_SIZE);
	cvAbsDiff(avg, imgSmooth, diff);
	cvConvertScale(diff, diff, 5);
	cvShowImage(BGEST, diff);
	save("bgestimate", diff);
	return diff;
}

/* Calculate average intensity of a region within an image */
double avg(IplImage *img, CvRect &r) {
	cvSetImageROI(img, r);
	CvScalar s = cvAvg(img);
	cvResetImageROI(img);
	return s.val[0];
}

int main(int argc, char *argv[]) {
	if (argc != 2) {
		printf("Usage: %s <directory>\n", argv[0]);
		exit(1);
	}

	cvNamedWindow(WINDOW, CV_WINDOW_AUTOSIZE);
	cvNamedWindow(BGEST, CV_WINDOW_AUTOSIZE);
	cvNamedWindow(MCONV, CV_WINDOW_AUTOSIZE);

	IplImage *imgA, *imgB, *imgResult, *imgTmp;

	imgTmp = loadImage(100, 0, argv[1]);
	CvSize imgSize = cvGetSize(imgTmp);
	cvMoveWindow(BGEST, imgSize.width, 0);
	cvMoveWindow(MCONV, imgSize.width*2, 0);
	imgA           = cvCreateImage(imgSize, IPL_DEPTH_8U, 1);
	imgB           = cvCreateImage(imgSize, IPL_DEPTH_8U, 1);
	IplConvKernel *kernel1  = cvCreateStructuringElementEx(SUPPORT,SUPPORT,SUPPORT/2,SUPPORT/2, CV_SHAPE_RECT);


	/* Our 3 rectangular areas where we will monitor for car traffic */
	CvRect leftLane  = cvRect(44,  163, 67, 73);
	CvRect rightLane = cvRect(166, 130, 50, 70);
	CvRect rightTurn = cvRect(227, 29,  30, 20);
	

	int down, up, right;
	down = up = right = 0;
	char buf[255];
	CvFont font;
	cvInitFont(&font, CV_FONT_HERSHEY_PLAIN, 1.0, 1.0, 0, 1, 8);

	int i = HISTORY+2;
	imgA = bgEstimate(i-1, imgSize, argv[1]);
	cvMorphologyEx(imgA, imgA, imgTmp, kernel1, CV_MOP_OPEN);

	/* Loop through all our images, our `movie' and perform a bgEstimate, and
	 * an opening. Next check for our regions if it contains enough intensity,
	 * if so a car is passing by.
	 */
	for (i=500; i <= 1664; i++) {
		imgB = bgEstimate(i, imgSize, argv[1]);
		cvMorphologyEx(imgB, imgB, imgTmp, kernel1, CV_MOP_OPEN);
		imgResult = loadImage(i, CV_LOAD_IMAGE_COLOR, argv[1]);
		cvRectangle(
			imgB, 
			cvPoint(leftLane.x, leftLane.y), 
			cvPoint(leftLane.x+leftLane.width, leftLane.y+leftLane.height), 
			CV_RGB(255,255,255),
			1
		);

		/* Check for the left lane */
		if (avg(imgB, leftLane) > 90 && avg(imgA, leftLane) < 90) {
			cvRectangle(
				imgResult, 
				cvPoint(leftLane.x, leftLane.y), 
				cvPoint(leftLane.x+leftLane.width, leftLane.y+leftLane.height), 
				CV_RGB(255,0,0),
				3
			);
			down++;
		}
		cvRectangle(
			imgB, 
			cvPoint(rightLane.x, rightLane.y), 
			cvPoint(rightLane.x+rightLane.width, rightLane.y+rightLane.height), 
			CV_RGB(255,255,255),
			1
		);

		/* Check for the right lane */
		if (avg(imgB, rightLane) > 130 && avg(imgA, rightLane) < 130) {
			cvRectangle(
				imgResult, 
				cvPoint(rightLane.x, rightLane.y), 
				cvPoint(rightLane.x+rightLane.width, rightLane.y+rightLane.height), 
				CV_RGB(0,0,255),
				3
			);
			up++;
		}
		cvRectangle(
			imgB, 
			cvPoint(rightTurn.x, rightTurn.y), 
			cvPoint(rightTurn.x+rightTurn.width, rightTurn.y+rightTurn.height), 
			CV_RGB(255,255,255),
			1
		);

		/* Check for right turns */
		if (avg(imgB, rightTurn) > 60 && avg(imgA, rightTurn) < 60) {
			cvRectangle(
				imgResult, 
				cvPoint(rightTurn.x, rightTurn.y), 
				cvPoint(rightTurn.x+rightTurn.width, rightTurn.y+rightTurn.height), 
				CV_RGB(0,255,0),
				3
			);
			right++;
		}

		/* Print output to screen */
		sprintf(buf, "Down %d, Up %d, Right %d", down, up, right);
		cvShowImage(MCONV, imgB);
		save("convolution", imgB);

		cvPutText(imgResult, buf, cvPoint(10,20), &font, CV_RGB(255,0,0));
		cvShowImage(WINDOW, imgResult);
		save("final", imgResult);

		cvWaitKey(41); /* 24 fps */
		/* save calculations */
		imgA = cvCloneImage(imgB);
	}

    return 0;
}
